import pandas as pd
import numpy as np
from pathlib import Path
import logging
from datetime import datetime


# 配置日志
def setup_logging():
    """设置日志配置"""
    # 创建logs文件夹（如果不存在）
    log_dir = Path('logs')
    log_dir.mkdir(exist_ok=True)

    # 生成日志文件名（包含时间戳）
    log_filename = f"data_cleaning_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
    log_path = log_dir / log_filename

    # 配置日志格式
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_path, encoding='utf-8'),
            logging.StreamHandler()
        ]
    )


def clean_athletes_data():
    """清理运动员数据文件"""
    try:
        logging.info("开始处理运动员数据...")
        file_path = Path('summerOly_athletes.csv')
        df = pd.read_csv(file_path)

        # 1. 处理缺失值
        df['Medal'] = df['Medal'].fillna('No medal')

        # 2. 标准化字符串列
        string_columns = ['Sex', 'NOC', 'Team', 'Sport', 'Event']
        for col in string_columns:
            if col in df.columns:
                df[col] = df[col].str.strip()

        # 3. 确保年份为整数
        df['Year'] = pd.to_numeric(df['Year'], errors='coerce')

        # 4. 删除完全重复的行
        df = df.drop_duplicates()

        # 保存更新后的数据
        df.to_csv(file_path, index=False)
        logging.info("运动员数据清理完成并保存")

        stats = {
            "总记录数": len(df),
            "缺失值统计": df.isnull().sum().to_dict(),
            "唯一国家数": len(df['NOC'].unique()),
            "运动项目数": len(df['Sport'].unique())
        }
        return stats

    except Exception as e:
        logging.error(f"处理运动员数据时出错: {e}", exc_info=True)
        raise


def clean_medal_counts():
    """清理奖牌统计数据文件"""
    try:
        logging.info("开始处理奖牌统计数据...")
        file_path = Path('summerOly_medal_counts.csv')
        df = pd.read_csv(file_path)

        # 1. 确保奖牌数为数值类型
        medal_columns = ['Gold', 'Silver', 'Bronze', 'Total']
        for col in medal_columns:
            df[col] = pd.to_numeric(df[col], errors='coerce')

        # 2. 验证并修正奖牌总数
        df['Calculated_Total'] = df['Gold'] + df['Silver'] + df['Bronze']
        mask = df['Total'] != df['Calculated_Total']
        if mask.any():
            logging.info(f"发现{mask.sum()}条奖牌总数不匹配记录，已更正")
            df.loc[mask, 'Total'] = df.loc[mask, 'Calculated_Total']

        df = df.drop('Calculated_Total', axis=1)

        # 3. 确保年份为整数
        df['Year'] = pd.to_numeric(df['Year'], errors='coerce')

        # 4. 标准化NOC代码
        df['NOC'] = df['NOC'].str.strip()

        # 5. 删除重复记录
        df = df.drop_duplicates()

        # 保存更新后的数据
        df.to_csv(file_path, index=False)
        logging.info("奖牌统计数据清理完成并保存")

        stats = {
            "总记录数": len(df),
            "年份范围": f"{df['Year'].min()}-{df['Year'].max()}",
            "总金牌数": df['Gold'].sum(),
            "总奖牌数": df['Total'].sum()
        }
        return stats

    except Exception as e:
        logging.error(f"处理奖牌统计数据时出错: {e}", exc_info=True)
        raise


def clean_hosts_data():
    """清理主办国数据文件"""
    try:
        logging.info("开始处理主办国数据...")
        file_path = Path('summerOly_hosts.csv')
        df = pd.read_csv(file_path)

        # 1. 处理取消的奥运会
        cancelled_mask = df['Host'].str.contains('Cancelled', na=False, case=False)
        df.loc[cancelled_mask, 'Host'] = 'Cancelled'

        # 2. 对非取消的记录进行城市和国家分离
        non_cancelled = ~cancelled_mask
        df.loc[non_cancelled, 'Host'] = df.loc[non_cancelled, 'Host'].str.strip()

        # 创建新的城市和国家列
        df['City'] = None
        df['Country'] = None

        # 对非取消的记录进行分离处理
        temp = df.loc[non_cancelled, 'Host'].str.split(',', expand=True)
        df.loc[non_cancelled, 'City'] = temp[0].str.strip()
        df.loc[non_cancelled, 'Country'] = temp[1].str.strip()

        # 对取消的记录填充特殊值
        df.loc[cancelled_mask, ['City', 'Country']] = 'Cancelled'

        # 3. 确保年份为整数
        df['Year'] = pd.to_numeric(df['Year'], errors='coerce')

        # 保存更新后的数据
        df.to_csv(file_path, index=False)
        logging.info("主办国数据清理完成并保存")

        stats = {
            "总记录数": len(df),
            "取消的奥运会数量": cancelled_mask.sum(),
            "主办国家数量": len(df.loc[~cancelled_mask, 'Country'].unique())
        }
        return stats

    except Exception as e:
        logging.error(f"处理主办国数据时出错: {e}", exc_info=True)
        raise


def clean_programs_data():
    """清理项目数据文件"""
    try:
        logging.info("开始处理项目数据...")
        file_path = Path('summerOly_programs.csv')

        # 尝试不同的编码方式读取文件
        encodings = ['utf-8', 'latin1', 'iso-8859-1', 'cp1252']
        df = None

        for encoding in encodings:
            try:
                df = pd.read_csv(file_path, encoding=encoding)
                logging.info(f"成功使用 {encoding} 编码读取文件")
                break
            except UnicodeDecodeError:
                logging.warning(f"使用 {encoding} 编码读取失败，尝试下一种编码")
                continue

        if df is None:
            raise ValueError("无法使用任何编码方式读取文件")

        # 1. 处理特殊字符
        df = df.replace('�', np.nan)

        # 2. 清理字符串列
        string_columns = ['Sport', 'Discipline']
        for col in string_columns:
            df[col] = df[col].str.strip()

        # 3. 处理年份列
        year_columns = [str(year) for year in range(1896, 2025, 4)]
        for col in year_columns:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
                df[col] = df[col].fillna(0)

        # 4. 删除全空的行
        df = df.dropna(subset=['Sport', 'Discipline'], how='all')

        # 保存更新后的数据（使用成功的编码方式）
        df.to_csv(file_path, index=False, encoding='utf-8')
        logging.info("项目数据清理完成并保存")

        stats = {
            "总记录数": len(df),
            "运动项目数": len(df['Sport'].unique()),
            "项目大项数": len(df['Discipline'].unique())
        }
        return stats

    except Exception as e:
        logging.error(f"处理项目数据时出错: {e}", exc_info=True)
        raise


def extract_non_medal_countries_features():
    """为从未获得过奖牌的国家提取运动员特征"""
    try:
        # 读取原始运动员数据和未获奖国家列表
        athletes_df = pd.read_csv('summerOly_athletes.csv')
        non_medal_countries = pd.read_csv('non_medal_countries.csv')
        
        # 创建特征存储字典
        features = {}
        
        # 只处理2000年之后的数据，使特征更具现代意义
        recent_athletes = athletes_df[athletes_df['Year'] >= 2000]
        
        for country in non_medal_countries['Country']:
            # 获取该国运动员数据
            country_athletes = recent_athletes[recent_athletes['Team'] == country]
            
            if len(country_athletes) == 0:
                continue
                
            # 1. 基础统计特征
            athlete_count = len(country_athletes)
            unique_sports = len(country_athletes['Sport'].unique())
            participation_years = len(country_athletes['Year'].unique())
            
            # 2. 时间趋势特征
            years = sorted(country_athletes['Year'].unique())
            yearly_counts = country_athletes.groupby('Year').size()
            if len(years) >= 2:
                trend = np.polyfit(years, yearly_counts, 1)[0]
            else:
                trend = 0
                
            # 3. 运动项目多样性
            sport_concentration = (
                country_athletes.groupby('Sport').size() / 
                len(country_athletes)
            ).max()
            
            # 4. 最近参与情况
            recent_years = [2020, 2024]
            recent_participation = len(
                country_athletes[
                    country_athletes['Year'].isin(recent_years)
                ]
            )
            
            # 5. 运动员经验
            athlete_appearances = country_athletes.groupby('Name').size()
            avg_experience = athlete_appearances.mean()
            max_experience = athlete_appearances.max()
            
            # 存储特征
            features[country] = {
                'country': country,
                'athlete_count': athlete_count,
                'unique_sports': unique_sports,
                'participation_years': participation_years,
                'athlete_trend': trend,
                'sport_concentration': sport_concentration,
                'recent_participation': recent_participation,
                'avg_experience': avg_experience,
                'max_experience': max_experience,
                'last_year': years[-1] if years else 0
            }
        
        # 转换为DataFrame
        features_df = pd.DataFrame.from_dict(features, orient='index')
        
        # 保存特征
        output_path = 'features/non_medal_countries_features.csv'
        features_df.to_csv(output_path, index=False)
        
        logging.info(f"未获奖国家特征提取完成，已保存到: {output_path}")
        
        # 返回统计信息
        stats = {
            "处理国家数": len(features),
            "特征数量": len(features_df.columns),
            "有效数据国家数": len(features_df[features_df['athlete_count'] > 0])
        }
        return stats
        
    except Exception as e:
        logging.error(f"未获奖国家特征提取失败: {str(e)}")
        raise


def main():
    """执行所有数据清理流程"""
    try:
        # 设置日志
        setup_logging()
        logging.info("开始数据清理流程...")

        # 创建统计信息字典
        stats = {}

        # 依次执行各个清理函数
        stats['athletes'] = clean_athletes_data()
        stats['medals'] = clean_medal_counts()
        stats['hosts'] = clean_hosts_data()
        stats['programs'] = clean_programs_data()
        stats['non_medal_countries'] = extract_non_medal_countries_features()
        print(stats)
        # 保存统计信息
        with open('data_cleaning_report.txt', 'w', encoding='utf-8') as f:
            for category, category_stats in stats.items():
                f.write(f"\n{category.upper()} 数据清理统计:\n")
                for key, value in category_stats.items():
                    f.write(f"{key}: {value}\n")

        logging.info("所有数据清理完成")

    except Exception as e:
        logging.error(f"数据清理过程中出现错误: {e}", exc_info=True)
        raise


if __name__ == "__main__":
    main()
